# -*- coding: utf-8 -*-
"""
@Time    : 2024/6/29 09:59 
@Author  : ZhangShenao 
@File    : 1.StdOutCallbackHandler使用.py 
@Desc    : 使用StdOutCallbackHandler回调机制,将整个chain生命周期的关键链路打印到控制台
"""
from langchain_community.chat_models import ChatZhipuAI
from langchain_core.callbacks import StdOutCallbackHandler
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough

from llm_ops_callback_handler import LLMOpsCallbackHandler

# 编排Prompt
prompt = ChatPromptTemplate.from_template('{query}')

# 构建llm
llm = ChatZhipuAI(model="GLM-4-Air")

# 创建OutputParser
parser = StrOutputParser()

# 构建chain
# 使用RunnablePassthrough简化调用
chain = {'query': RunnablePassthrough()} | prompt | llm | parser

# 以stream模式运行chain
content = chain.stream('你好，你是什么模型？',
                       config={
                           # 注册StdOutCallbackHandler内置回调和LLMOpsCallbackHandler自定义回调
                           'callbacks': [StdOutCallbackHandler(), LLMOpsCallbackHandler()],
                       })

for chunk in content:
    # print(chunk, end='')
    len(chunk)

print('\n')
